'''
读取csv格式文件，提取前3列，然后增加5列数据，其中前4以0.0填充，最后一列以1.0填充，最后以csv格式保存成新的文件。
采用命令行参数传入输入、输出的文件名
'''
import pandas as pd
import sys
import argparse

# 创建ArgumentParser对象
parser = argparse.ArgumentParser()

parser.add_argument('input_file', type=str, help='Enter the input file path')
parser.add_argument('output_file', type=str, help='Enter the output file path')

parser.add_argument('--use_fused', type=int, default=0, 
                    help='use rtk vio fused trajectory(1) or vio trajectory(0), default 1')

parser.add_argument('--skip_rows', type=int, default=1, 
                    help='skip the first n rows, until vio status(the fifth column) is equal or greater '
                    'than 4, default 1')

parser.add_argument('--truncate', type=int, default=0, 
                    help='truncate rows for row index > truncate, default not truncating')

# 解析命令行参数
args = parser.parse_args()

# read csv file
# input_file = sys.argv[1]
# output_file = sys.argv[2]
data = pd.read_csv(args.input_file, sep=' ',header=None, comment='#')

if args.skip_rows:
    print('skip the first n rows, until vio status(the fifth column) is equal or greater than 4')
    skipped_lines = 0
    # 迭代数据的每一行
    for index, row in data.iterrows():
        # 判断每行第五列数据是否大于4
        if row.iloc[5] >= 4:
            print(f"Row {index}: Fifth column value is equal or greater than 4, skip the first {index} rows.")
            skipped_lines =index
            break

    # 跳过前skipped_lines行
    data = data.iloc[skipped_lines:, :]

# truncate rows for large file
if args.truncate:
    max_row = args.truncate
    print(f'truncate rows for row index >= {max_row}.')
    data = data.iloc[:max_row, :]

# get the first three columns(timestamp, x, y)
new_data = []
if args.use_fused:
    print('use rtk vio fused trajectory') 
    new_data = data.iloc[:, :3]
else:
    print('use vio trajectory')
    new_data = data.iloc[:, 6:9]

# 给数据列增加header
new_header = ['timestamp', 'x', 'y']
new_data.columns = new_header

# add columns
new_data['z'] = 0.0
new_data['qx'] = 0.0
new_data['qy'] = 0.0
new_data['qz'] = 0.0
new_data['qw'] = 1.0

print(new_data.head())

print("save data to file: ", args.output_file)
new_data.to_csv(args.output_file, sep=' ', header=False, index=False)

